#!/usr/bin/env python
# Copyright 2014-2016 RethinkDB, all rights reserved.

# issue-url: https://github.com/rethinkdb/rethinkdb/issues/3038
# issue-description: when a primary replica suddenly dies changefeed requests to that server are not cleaned up correctly

import os, sys

sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), os.path.pardir, 'common'))
import driver, utils

r = utils.import_python_driver()
dbName, tableName = utils.get_test_db_table()

# == constants

records = 5000

# == start up two joined servers

with driver.Cluster(initial_servers=2, console_output=True) as cluster:
    
    stableServer = cluster[0]
    sacrificeServer = cluster[1]
    
    # == create a table
    
    conn = r.connect(host=stableServer.host, port=stableServer.driver_port)
    
    if dbName not in r.db_list().run(conn):
        r.db_create(dbName).run(conn)
    if tableName in r.db(dbName).table_list().run(conn):
        r.db(dbName).table_drop(tableName).run(conn)
    r.db(dbName).table_create(tableName).run(conn)
    
    # == populate the table
    
    r.db(dbName).table(tableName).insert(r.range(1, records + 1).map({'id':r.row})).run(conn)
    assert r.db(dbName).table(tableName).count().run(conn) == records
    
    # == shard the table
    
    shardPlan = [
        {'primary_replica':stableServer.name, 'replicas':[stableServer.name]},
        {'primary_replica':sacrificeServer.name, 'replicas':[sacrificeServer.name]}
    ]
    assert (r.db(dbName).table(tableName).config().update({'shards':shardPlan}).run(conn))['errors'] == 0
    r.db(dbName).table(tableName).wait(wait_for="all_replicas_ready").run(conn)
    
    # == open a changefeed
    
    changefeed = r.db(dbName).table(tableName).changes().run(conn)
    
    # == kill the second server
    
    sacrificeServer.kill()
    
    # == pull something from the changefeed
    
    try:
        next(changefeed)
        sys.exit('Failure: did not get a exception on the changefeed as expected!')
    except r.errors.ReqlRuntimeError:
        pass
    
    # == check that the first server can close gracefully
    
    try:
        stableServer.check_and_stop()
    except RuntimeError as e:
        sys.stderr.write('Failure: server did not close cleanly: %s\n' % str(e))
        print('========= Stdout/Stderr: =========')
        sys.stdout.write(stableServer.console_output)
        print('\n=========')
        sys.exit(1)

print('Completed Successfully')
